#include "printers.h"
#include "file_cmd.h"

/****************************************************************************\
* Private Variables/Functions
\****************************************************************************/

#define S_MAX_METHOD_NAME_LEN 300

void s_get_method_name(str_w str_buf, size_t str_len, pe_file *p_pe, size_t rid_method)
{
	//
	// TODO: Handle nested types....-_-
	//
	mdt_tbl_def *p_tbl_method = p_pe->metadata.tbl_defs + TBLID_Method;
	mdt_tbl_def *p_tbl_type = p_pe->metadata.tbl_defs + TBLID_TypeDef;

	mdt_col_def *p_col_method_name = p_pe->metadata.col_defs + p_tbl_method->cols + 3;
	mdt_col_def *p_col_type_name = p_pe->metadata.col_defs + p_tbl_type->cols + 1;
	mdt_col_def *p_col_type_ns = p_pe->metadata.col_defs + p_tbl_type->cols + 2;
	mdt_col_def *p_col_type_method = p_pe->metadata.col_defs + p_tbl_type->cols + 5;

	byte *p_type_name, *p_type_ns, *p_method_name;
	byte *p_row_data, *p_col_data;
	size_t ridz_type, type_method_rid, next_type_method_rid;
	u4 method_name_value, type_name_value, type_ns_value;

	size_t n_write = 0;

	// Find id of type that contains the method
	for(ridz_type = 0; ridz_type < p_tbl_type->record_num; ++ridz_type)
	{
		p_row_data = p_tbl_type->p_data + ridz_type * p_tbl_type->record_size;
		p_col_data = p_row_data + p_col_type_method->offset;

		if(p_col_type_method->size == 2)
		{
			type_method_rid = *((u2 *)p_col_data);
			next_type_method_rid = ridz_type < p_tbl_type->record_num - 1
				? *((u2 *)(p_col_data + p_tbl_type->record_size))
				: p_tbl_method->record_num;
		}
		else
		{
			type_method_rid = *((u4 *)p_col_data);
			next_type_method_rid = ridz_type < p_tbl_type->record_num - 1
				? *((u4 *)(p_col_data + p_tbl_type->record_size))
				: p_tbl_method->record_num;
		}

		if(type_method_rid <= rid_method && rid_method < next_type_method_rid)
			break;
	}

	// Find names
	p_col_data = p_row_data + p_col_type_name->offset;
	type_name_value = p_col_type_name->size == 2 ? *((u2 *)p_col_data) : *((u4 *)p_col_data);

	p_col_data = p_row_data + p_col_type_ns->offset;
	type_ns_value = p_col_type_ns->size == 2 ? *((u2 *)p_col_data) : *((u4 *)p_col_data);

	p_row_data = p_tbl_method->p_data + (rid_method - 1) * p_tbl_method->record_size;
	p_col_data = p_row_data + p_col_method_name->offset;
	method_name_value = p_col_method_name->size == 2 ? *((u2 *)p_col_data) : *((u4 *)p_col_data);

	// Find #Strings data.
	p_type_name   = p_pe->p_base + p_pe->metadata.strings_offset + type_name_value;
	p_type_ns     = p_pe->p_base + p_pe->metadata.strings_offset + type_ns_value;
	p_method_name = p_pe->p_base + p_pe->metadata.strings_offset + method_name_value;

	while(n_write < str_len && *p_method_name)
		str_buf[n_write++] = *p_method_name++;

	if(n_write < str_len) str_buf[n_write++] = L',';
	if(n_write < str_len) str_buf[n_write++] = L' ';

	while(n_write < str_len && *p_type_name)
		str_buf[n_write++] = *p_type_name++;

	if(n_write < str_len) str_buf[n_write++] = L',';
	if(n_write < str_len) str_buf[n_write++] = L' ';

	while(n_write < str_len && *p_type_ns)
		str_buf[n_write++] = *p_type_ns++;
}

BOOL s_print_method_info(pe_file *p_pe, size_t rid_method)
{
	mdt_tbl_def *p_tbl;
	byte *p_row;
	size_t rva, offset;
	wchar_t method_name[S_MAX_METHOD_NAME_LEN + 1] = {0};

	if(!check_file_opened() || !check_pe_headers(p_pe) || !check_metadata(p_pe))
		return FALSE;

	p_tbl = p_pe->metadata.tbl_defs + TBLID_Method;

	if(rid_method < 1 || rid_method > p_tbl->record_num)
	{
		wprintf(L"<rid> is out of range (valid value should be 1-%d).\n",
			p_tbl->record_num);
		return FALSE;
	}

	p_row = p_tbl->p_data + p_tbl->record_size * (rid_method - 1);
	s_get_method_name(method_name, S_MAX_METHOD_NAME_LEN, p_pe, rid_method);

	wprintf(L"Method[%Xh (%d)]: %s\n", rid_method, rid_method, method_name);

	// RVA is the first column of Method table.
	rva = *((u4 *)p_row);

	if(!rva)
	{
		wprintf(L"This method has no IL codes.\n");
		return FALSE;
	}

	offset = rva_to_offset(p_pe, rva);
	wprintf(L"RVA=%08Xh, FileOffset=%08Xh\n", rva, offset);

	return TRUE;
}

void s_print_method_flags(cor_il_method_flags flags)
{
	switch(flags & CorILMethod_FormatMask)
	{
	case CorILMethod_TinyFormat : wprintf(L" TinyFormat");  break;
	case CorILMethod_SmallFormat: wprintf(L" SmallFormat"); break;
	case CorILMethod_FatFormat  : wprintf(L" FatFormat");   break;
	case CorILMethod_TinyFormat1: wprintf(L" TinyFormat");  break;
	}

	if(flags & CorILMethod_InitLocals)   wprintf(L" InitLocals");
	if(flags & CorILMethod_MoreSects)    wprintf(L" MoreSects");
	if(flags & CorILMethod_CompressedIL) wprintf(L" CompressedIL");

	wprintf(L"\n");
}

void s_print_ilmethod_tiny(pe_file *p_pe, IMAGE_COR_ILMETHOD_TINY *ph)
{
	wprintf(L"IMAGE_COR_ILMETHOD_TINY:\n");
	wprintf(L"    %08Xh %02Xh Flags_CodeSize %04Xh\n",
		PTR_D(p_pe->p_base, &ph->Flags_CodeSize),
		sizeof(ph->Flags_CodeSize),
		ph->Flags_CodeSize);

	wprintf(L"        Flags: %Xh ", ph->Flags);
	s_print_method_flags((cor_il_method_flags)ph->Flags);
	wprintf(L"        CodeSize: %Xh (%d)\n", ph->CodeSize, ph->CodeSize);
}

void s_print_ilmethod_fat(pe_file *p_pe, IMAGE_COR_ILMETHOD_FAT *ph)
{
	wprintf(L"IMAGE_COR_ILMETHOD_FAT:\n");
	wprintf(L"    %08Xh %02Xh Flags_Size_MaxStack %08Xh\n",
		PTR_D(p_pe->p_base, &ph->Flags_Size_MaxStack),
		sizeof(ph->Flags_Size_MaxStack),
		ph->Flags_Size_MaxStack);

	wprintf(L"        Flags: %Xh ", ph->Flags);
	s_print_method_flags((cor_il_method_flags)ph->Flags);
	wprintf(L"        Size: %Xh (%d)\n", ph->Size, ph->Size);
	wprintf(L"        MaxStack: %Xh (%d)\n", ph->MaxStack, ph->MaxStack);

	wprintf(L"    %08Xh %02Xh            CodeSize %08Xh (%d)\n",
		PTR_D(p_pe->p_base, &ph->CodeSize), sizeof(ph->CodeSize),
		ph->CodeSize, ph->CodeSize);
	wprintf(L"    %08Xh %02Xh      LocalVarSigTok %08Xh (%d)\n",
		PTR_D(p_pe->p_base, &ph->LocalVarSigTok), sizeof(ph->LocalVarSigTok),
		ph->LocalVarSigTok, ph->LocalVarSigTok);
}

void s_print_method_hdr(pe_file *p_pe, il_method *p_m)
{
	if(COR_ILMETHOD_IS_FAT(p_m->p_hdr))
		s_print_ilmethod_fat(p_pe, &p_m->p_hdr->Fat);
	else
		s_print_ilmethod_tiny(p_pe, &p_m->p_hdr->Tiny);
}

void s_print_method_sect_eh_small(pe_file *p_pe, IMAGE_COR_ILMETHOD_SECT_EH_SMALL *p_sect)
{
	size_t clause_num = (p_sect->SectSmall.DataSize
		- (sizeof(IMAGE_COR_ILMETHOD_SECT_EH_SMALL) - sizeof(IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_SMALL)))
		/ sizeof(IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_SMALL);
	size_t i, base_offset;
	IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_SMALL *p_clause;
	wchar_t token_str[METADATA_TOKEN_STR_LEN + 1];

	wprintf(L"IMAGE_COR_ILMETHOD_SECT_EH_SMALL (%d clauses):\n", clause_num);

	wprintf(L"    %08Xh %02Xh     Kind %02Xh\n",
		PTR_D(p_pe->p_base, &p_sect->SectSmall.Kind), sizeof(p_sect->SectSmall.Kind),
		p_sect->SectSmall.Kind);
	wprintf(L"    %08Xh %02Xh DataSize %02Xh (%d)\n",
		PTR_D(p_pe->p_base, &p_sect->SectSmall.DataSize), sizeof(p_sect->SectSmall.DataSize),
		p_sect->SectSmall.DataSize, p_sect->SectSmall.DataSize);

	for(i = 0, p_clause = p_sect->Clauses;
		i < clause_num;
		++i, ++p_clause)
	{
		base_offset = PTR_D(p_pe->p_base, p_clause);

		wprintf(L"\nClauses[%d] (IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_SMALL):\n", i + 1);
		wprintf(L"    %08Xh %02Xh         Flags %04Xh ",
			base_offset, 2,
			p_clause->Flags);
		switch(p_clause->Flags)
		{
		case COR_ILEXCEPTION_CLAUSE_NONE:       wprintf(L"(COR_ILEXCEPTION_CLAUSE_NONE)\n");       break;
		case COR_ILEXCEPTION_CLAUSE_FILTER:     wprintf(L"(COR_ILEXCEPTION_CLAUSE_FILTER)\n");     break;
		case COR_ILEXCEPTION_CLAUSE_FINALLY:    wprintf(L"(COR_ILEXCEPTION_CLAUSE_FINALLY)\n");    break;
		case COR_ILEXCEPTION_CLAUSE_FAULT:      wprintf(L"(COR_ILEXCEPTION_CLAUSE_FAULT)\n");      break;
		case COR_ILEXCEPTION_CLAUSE_DUPLICATED: wprintf(L"(COR_ILEXCEPTION_CLAUSE_DUPLICATED)\n"); break;
		default: wprintf(L"\n"); break;
		}
		wprintf(L"    %08Xh %02Xh     TryOffset %04Xh     (From IL_%08X)\n",
			base_offset += 2, 2,
			p_clause->TryOffset, p_clause->TryOffset);
		wprintf(L"    %08Xh %02Xh     TryLength %02Xh       (  To IL_%08X)\n",
			base_offset += 2, 1,
			p_clause->TryLength, p_clause->TryOffset + p_clause->TryLength);
		wprintf(L"    %08Xh %02Xh HandlerOffset %04Xh     (From IL_%08X)\n",
			base_offset += 1, 2,
			p_clause->HandlerOffset, p_clause->HandlerOffset);
		wprintf(L"    %08Xh %02Xh HandlerLength %02Xh       (  To IL_%08X)\n",
			base_offset += 2, 1,
			p_clause->HandlerLength, p_clause->HandlerOffset + p_clause->HandlerLength);

		switch(p_clause->Flags)
		{
		case COR_ILEXCEPTION_CLAUSE_NONE:
			format_metadata_token(p_pe, p_clause->ClassToken, token_str, METADATA_TOKEN_STR_LEN);
			wprintf(L"    %08Xh %02Xh    ClassToken %08Xh (%s)\n",
				PTR_D(p_pe->p_base, &p_clause->ClassToken), sizeof(p_clause->ClassToken),
				p_clause->ClassToken, token_str);
			break;
		case COR_ILEXCEPTION_CLAUSE_FILTER:
			wprintf(L"    %08Xh %02Xh  FilterOffset %08Xh (%d)\n",
				PTR_D(p_pe->p_base, &p_clause->FilterOffset), sizeof(p_clause->FilterOffset),
				p_clause->FilterOffset, p_clause->FilterOffset);
			break;
		default:
			wprintf(L"    %08Xh %02Xh     <Ignored> %08Xh\n",
				PTR_D(p_pe->p_base, &p_clause->ClassToken), sizeof(p_clause->ClassToken),
				p_clause->ClassToken);
		}
	}
}

void s_print_method_sect_eh_fat(pe_file *p_pe, IMAGE_COR_ILMETHOD_SECT_EH_FAT *p_sect)
{
	size_t clause_num = (p_sect->SectFat.DataSize
		- (sizeof(IMAGE_COR_ILMETHOD_SECT_EH_FAT) - sizeof(IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_FAT)))
		/ sizeof(IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_FAT);
	size_t i, base_offset;
	IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_FAT *p_clause;
	wchar_t token_str[METADATA_TOKEN_STR_LEN + 1];

	wprintf(L"IMAGE_COR_ILMETHOD_SECT_EH_FAT (%d clauses):\n", clause_num);

	base_offset = PTR_D(p_pe->p_base, p_sect);
	wprintf(L"    %08Xh %02Xh     Kind %02Xh\n",
		base_offset, 1,
		p_sect->SectFat.Kind);
	wprintf(L"    %08Xh %02Xh DataSize %06Xh (%d)\n",
		base_offset += 1, 3,
		p_sect->SectFat.DataSize, p_sect->SectFat.DataSize);

	for(i = 0, p_clause = p_sect->Clauses;
		i < clause_num;
		++i, ++p_clause)
	{
		wprintf(L"\nClauses[%d] (IMAGE_COR_ILMETHOD_SECT_EH_CLAUSE_FAT):\n", i + 1);
		wprintf(L"    %08Xh %02Xh         Flags %08Xh ",
			PTR_D(p_pe->p_base, &p_clause->Flags), sizeof(p_clause->Flags),
			p_clause->Flags);
		switch(p_clause->Flags)
		{
		case COR_ILEXCEPTION_CLAUSE_NONE:       wprintf(L"(COR_ILEXCEPTION_CLAUSE_NONE)\n");       break;
		case COR_ILEXCEPTION_CLAUSE_FILTER:     wprintf(L"(COR_ILEXCEPTION_CLAUSE_FILTER)\n");     break;
		case COR_ILEXCEPTION_CLAUSE_FINALLY:    wprintf(L"(COR_ILEXCEPTION_CLAUSE_FINALLY)\n");    break;
		case COR_ILEXCEPTION_CLAUSE_FAULT:      wprintf(L"(COR_ILEXCEPTION_CLAUSE_FAULT)\n");      break;
		case COR_ILEXCEPTION_CLAUSE_DUPLICATED: wprintf(L"(COR_ILEXCEPTION_CLAUSE_DUPLICATED)\n"); break;
		default: wprintf(L"\n"); break;
		}
		wprintf(L"    %08Xh %02Xh     TryOffset %08Xh (From IL_%08X)\n",
			PTR_D(p_pe->p_base, &p_clause->TryOffset), sizeof(p_clause->TryOffset),
			p_clause->TryOffset, p_clause->TryOffset);
		wprintf(L"    %08Xh %02Xh     TryLength %08Xh (  To IL_%08X)\n",
			PTR_D(p_pe->p_base, &p_clause->TryLength), sizeof(p_clause->TryLength),
			p_clause->TryLength, p_clause->TryOffset + p_clause->TryLength);
		wprintf(L"    %08Xh %02Xh HandlerOffset %08Xh (From IL_%08X)\n",
			PTR_D(p_pe->p_base, &p_clause->HandlerOffset), sizeof(p_clause->HandlerOffset),
			p_clause->HandlerOffset, p_clause->HandlerOffset);
		wprintf(L"    %08Xh %02Xh HandlerLength %08Xh (  To IL_%08X)\n",
			PTR_D(p_pe->p_base, &p_clause->HandlerLength), sizeof(p_clause->HandlerLength),
			p_clause->HandlerLength, p_clause->HandlerOffset + p_clause->HandlerLength);

		switch(p_clause->Flags)
		{
		case COR_ILEXCEPTION_CLAUSE_NONE:
			format_metadata_token(p_pe, p_clause->ClassToken, token_str, METADATA_TOKEN_STR_LEN);
			wprintf(L"    %08Xh %02Xh    ClassToken %08Xh (%s)\n",
				PTR_D(p_pe->p_base, &p_clause->ClassToken), sizeof(p_clause->ClassToken),
				p_clause->ClassToken, token_str);
			break;
		case COR_ILEXCEPTION_CLAUSE_FILTER:
			wprintf(L"    %08Xh %02Xh  FilterOffset %08Xh (%d)\n",
				PTR_D(p_pe->p_base, &p_clause->FilterOffset), sizeof(p_clause->FilterOffset),
				p_clause->FilterOffset, p_clause->FilterOffset);
			break;
		default:
			wprintf(L"    %08Xh %02Xh     <Ignored> %08Xh\n",
				PTR_D(p_pe->p_base, &p_clause->ClassToken), sizeof(p_clause->ClassToken),
				p_clause->ClassToken);
		}
	}
}

void s_print_method_sect(pe_file *p_pe, byte *p_sect)
{
	u1 kind = *p_sect;

	switch(kind & CorILMethod_Sect_KindMask)
	{
	case CorILMethod_Sect_EHTable:
		if(kind & CorILMethod_Sect_FatFormat)
			s_print_method_sect_eh_fat(p_pe, (IMAGE_COR_ILMETHOD_SECT_EH_FAT *)p_sect);
		else
			s_print_method_sect_eh_small(p_pe, (IMAGE_COR_ILMETHOD_SECT_EH_SMALL *)p_sect);
		break;

	default:
		wprintf(L"Unknown section: kind=%Xh (%d).\n", kind, kind);
		break;
	}
}

void s_print_method_sects(pe_file *p_pe, il_method *p_m)
{
	byte *p_sect;

	wprintf(L"Extra sections:\n");

	for(p_sect = p_m->p_sect; p_sect; p_sect = next_il_method_sect(p_pe, p_sect))
	{
		if(p_sect > p_m->p_sect)
			wprintf(L"\n");

		s_print_method_sect(p_pe, p_sect);
	}
}

/****************************************************************************\
* Public Functions
\****************************************************************************/

void print_il_method(pe_file *p_pe, size_t rid_method, BOOL with_code)
{
	il_method ilm;
	size_t code_size;

	if(!s_print_method_info(p_pe, rid_method))
		return;

	if(!get_il_method(p_pe, rid_method, &ilm))
		return;

	wprintf(L"\n");
	s_print_method_hdr(p_pe, &ilm);
	wprintf(L"\n");

	code_size = COR_ILMETHOD_IS_FAT(ilm.p_hdr)
		? ilm.p_hdr->Fat.CodeSize : ilm.p_hdr->Tiny.CodeSize;

	wprintf(L"Method body: %08Xh-%08Xh, %Xh (%d) bytes.\n",
		PTR_D(p_pe->p_base, ilm.p_body),
		PTR_D(p_pe->p_base, ilm.p_body + code_size - 1),
		code_size, code_size);

	if(with_code)
	{
		wprintf(L"\n");
		print_il_stream(ilm.p_body, code_size);
	}

	if(COR_ILMETHOD_HAS_MORE_SECT(ilm.p_hdr))
	{
		wprintf(L"\n");
		s_print_method_sects(p_pe, &ilm);
	}
}
